"""
Authors:    Jingjing WU (吴京京) <https://github.com/wj-Mcat>

2020-now @ Jingjing Wu
Licensed under the Apache License, Version 2.0 (the 'License');
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an 'AS IS' BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
from __future__ import annotations
from typing import List, Union, Dict, Any
from dataclasses import dataclass, field
from dataclasses_json import dataclass_json

from torch import Tensor


@dataclass_json
@dataclass
class InputExample:
    """base class for all of input example"""
    input_id: Union[str, int]
    metadata: Dict[str, Any]


@dataclass_json
@dataclass
class TextClassificationInputExample(InputExample):
    """single text classification input example"""
    raw_text: str
    category: str = field(default_factory=str)

    # this field has value only when at inference stage
    score: float = field(default_factory=float)


@dataclass_json
@dataclass
class TextClassificationsInputExample(InputExample):
    """multi text classifications input example"""
    raw_text: str
    categories: List[str] = field(default_factory=list)

    # this field has value only when at inference stage
    scores: List[float] = field(default_factory=list)


@dataclass_json
@dataclass
class SequenceTaggingInputExample(InputExample):
    """sequence tagging input example

    tag: the tag of one token in sentence
    """
    raw_text: str
    tags: List[str] = field(default_factory=list)

    # this field has value only when at inference stage
    scores: List[List[float]] = field(default_factory=list)


@dataclass_json
@dataclass
class SentenceSimilarityInputExample(InputExample):
    """sentence similarity input example"""
    first_sentence: str
    second_sentence: str

    # this field has value only when at inference stage
    scores: List[float] = field(default_factory=list)
